from typing import Optional
import mindspore
import mindspore.ops as ops
def bincount(
    input: mindspore.Tensor, batch: Optional[mindspore.Tensor] = None, minlength: int = 0
):
    assert input.ndim == 1
    if batch is None:
        return ops.Bincount()(input, minlength)
    else:
        assert batch.shape == input.shape

        length = input.max().item() + 1
        if minlength == 0:
            minlength = length
        if length > minlength:
            raise ValueError(
                f"minlength {minlength} too small for input with integers up to and including {length}"
            )

        # Flatten indexes
        # Make each "class" in input into a per-input class.
        input = input + batch * minlength

        num_batch = batch.max() + 1

        return ops.Bincount()(input, minlength * num_batch).reshape(
            num_batch, minlength
        )
